
import torch.nn as nn

# 单网络loss
class Loss_one(nn.Module):
    
    def __init__(self,use_location):
        super(Loss_one, self).__init__()
        self.loss_cross = nn.CrossEntropyLoss()
        self.use_location = use_location

    # loss_one = 1*location的均方差 + Y的均方差
    def forward(self,Y_hat  ,Y  ,location_hat  ,location):

        loss_Y  = self.loss_cross(Y_hat,Y.reshape(-1))
        loss_one = loss_Y

        loss_location  = self.loss_cross(location_hat,location.reshape(-1))
        if self.use_location:
            loss_one+=loss_location

        return loss_one